import os

import cv2
import face_recognition
import numpy as np


def load_face_encodings(data_dir):
    """
    加载指定目录下的人脸图像并生成特征编码

    参数:
        data_dir: 存储人脸图像的目录路径

    返回:
        encodings: 字典，键为姓名，值为人脸特征编码
    """
    encodings = {}
    for root, dirs, files in os.walk(data_dir):
        for person in files:
            image_path = os.path.join(root, person)
            try:
                # 加载图像文件
                image = face_recognition.load_image_file(image_path)
                # 生成人脸特征编码
                face_encodings_list = face_recognition.face_encodings(image)

                if not face_encodings_list:
                    print(f"警告: 在 {image_path} 中未检测到人脸，已跳过")
                    continue

                # 取第一个检测到的人脸
                encoding = face_encodings_list
                if not encoding:
                    print(f"警告: {image_path} 中未检测到人脸，已跳过")
                    continue
                # 提取文件名中的姓名作为键（去除文件扩展名）
                encodings[person.split('.')[0]] = encoding[0]
            except Exception as e:
                print(f"处理 {image_path} 时出错: {str(e)}")

    return encodings


def mark_attendance(name):
    """
    记录考勤信息到CSV文件

    参数:
        name: 被识别到的人员姓名
    """
    # 以追加模式打开考勤记录文件
    # with open('attendance.csv', 'a') as f:
    #     # 获取当前时间并格式化
    #     current_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    #     # 写入姓名和打卡时间
    #     f.write(f'{name},{current_time}\n')
    print(f'{name}已打卡成功！')


# 人脸图像数据目录
def_data_dir = 'attendance_photos'
# 加载所有已知人脸的特征编码
face_encodings = load_face_encodings(def_data_dir)

def search_face():
    # 初始化摄像头（0表示默认摄像头）
    cap = cv2.VideoCapture(0)

    # 实时人脸检测与识别循环
    while True:
        # 读取摄像头帧
        ret, frame = cap.read()
        # 将帧缩小为原来的1/4以提高处理速度
        small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
        # 转换颜色空间从BGR（OpenCV默认）到RGB（face_recognition库需要）
        rgb_small_frame = np.ascontiguousarray(small_frame[:, :, ::-1])

        # 检测当前帧中的所有人脸位置
        face_locations = face_recognition.face_locations(rgb_small_frame)
        # 生成当前帧中所有人脸的特征编码
        face_encodings_frame = face_recognition.face_encodings(rgb_small_frame, face_locations)

        # 遍历当前帧中的每个人脸
        for (top, right, bottom, left), face_encoding in zip(face_locations, face_encodings_frame):
            # 与已知人脸编码进行比对
            matches = face_recognition.compare_faces(list(face_encodings.values()), face_encoding)
            name = "Unknown"  # 默认未知身份

            # 如果找到匹配的已知人脸
            if True in matches:
                first_match_index = matches.index(True)
                name = list(face_encodings.keys())[first_match_index]
                # 记录考勤（可以添加逻辑避免重复打卡）
                mark_attendance(name)

            # 将人脸位置从缩小的帧转换回原始帧的坐标（因为之前缩小了4倍）
            top *= 4
            right *= 4
            bottom *= 4
            left *= 4

            # 绘制绿色边框（BGR格式，(0, 255, 0)表示绿色）
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)

            # 在边框上方绘制姓名标签
            # 绘制标签背景
            cv2.rectangle(frame, (left, top - 35), (right, top), (0, 255, 0), cv2.FILLED)
            # 设置字体
            font = cv2.FONT_HERSHEY_DUPLEX
            # 绘制姓名
            # cv2.puttext不支持中文,通过Pillow转换成位图添加上去 参考MTCNN
            cv2.putText(frame, name, (left + 6, top - 6), font, 1.0, (255, 255, 255), 1)

        # 显示视频帧窗口
        cv2.imshow('Face Attendance System', frame)

        # 按'q'键退出程序
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放摄像头资源
    cap.release()
    # 关闭所有OpenCV窗口
    cv2.destroyAllWindows()

if __name__ == '__main__':
    search_face()